from config import *
import os
import requests
import re
import json
import time
import threading
from prometheus_api_client import PrometheusConnect
from flask import Flask, request, jsonify, render_template
from datetime import datetime
from wechat_service import wechat_service
from wechatpy.enterprise.crypto import WeChatCrypto
from wechatpy.exceptions import InvalidSignatureException

# 初始化flask
app = Flask(__name__)
app.json.ensure_ascii = False

def create_wechat_crypto():
    """创建企业微信加密器"""
    token = os.getenv('WECHAT_TOKEN', '')
    encoding_aes_key = os.getenv('WECHAT_ENCODING_AES_KEY', '')
    corp_id = os.getenv('WECHAT_CORP_ID', '')
    
    if token and encoding_aes_key and corp_id:
        return WeChatCrypto(token, encoding_aes_key, corp_id)
    return None

def generate_help_message():
    """生成统一的帮助信息"""
    projects_str = ', '.join(projects)
    query_type_str = ', '.join(promql_list)
    
    help_msg = (
        f'🤖 {BOT_NAME} 帮助信息\n\n'
        f'📋 当前支持项目: {projects_str}\n'
        f'📊 当前支持查询项: {query_type_str}\n\n'
        f'💡 使用方法:\n'
        f'• 查询项目需要和查询事项空格隔开\n'
        f'• 支持模糊查询，如: "联投 内存" 或 "项目名 CPU利用率"\n'
        f'• 示例: "联投 内存使用率"'
    )
    return help_msg

def cleanup_connections_periodically():
    """定期清理过期的Prometheus连接"""
    while True:
        try:
            time.sleep(30)  # 每30秒检查一次
            wechat_service.cleanup_prometheus_connection()
        except Exception as e:
            logger.error(f"定期清理连接异常: {e}")

# 添加请求日志中间件
@app.before_request
def log_request_info():
    """记录请求信息"""
    access_logger.info(f"请求开始 - {request.method} {request.path} - 来源IP: {request.remote_addr}")
    access_logger.info(f"请求头: {dict(request.headers)}")
    if request.is_json:
        access_logger.info(f"请求体: {request.get_json()}")

@app.after_request
def log_response_info(response):
    """记录响应信息"""
    access_logger.info(f"响应状态: {response.status_code}")
    access_logger.info(f"响应头: {dict(response.headers)}")
    return response

# 获取prometheus监控中所有的project
try:
    projects_response = requests.get(url=f'{PROMETHEUS_URL}/api/v1/label/project_name/values', timeout=10)
    projects_response.raise_for_status()
    projects = projects_response.json()['data']
    logger.info(f"成功获取Prometheus项目列表: {len(projects)}个项目")
    logger.info(f"Prometheus项目列表: {projects}")
except Exception as e:
    logger.error(f"获取Prometheus项目列表失败: {e}")
    projects = []

promql_list = list(PARAM_TO_PROMQL.keys())
def findproject_and_promql(inputstr):
    """查找匹配的项目和PromQL查询类型"""
    logger.info(f"开始解析用户输入: '{inputstr}'")
    
    keywords = inputstr.split()
    matched_project = None
    matched_promql = None
    
    logger.debug(f"提取的关键词: {keywords}")
    logger.debug(f"可用项目列表: {projects}")
    logger.debug(f"可用查询类型: {promql_list}")
    
    # 匹配项目
    for keyword in keywords:
        for project in projects:
            if re.search(keyword, project, re.IGNORECASE):
                matched_project = project
                logger.info(f"匹配到项目: '{project}' (关键词: '{keyword}')")
                break
        if matched_project:
            break

    # 匹配 PromQL
    for keyword in keywords:
        for query in promql_list:
            if re.search(keyword, query, re.IGNORECASE):
                matched_promql = PARAM_TO_PROMQL[query]
                logger.info(f"匹配到查询类型: '{query}' -> '{matched_promql}' (关键词: '{keyword}')")
                break
        if matched_promql:
            break

    logger.info(f"解析结果 - 项目: {matched_project}, 查询类型: {matched_promql}")
    return matched_project, matched_promql

def build_promql(project_name, query_type):
    """构建并执行PromQL查询"""
    start_time = time.time()
    logger.info(f"开始执行PromQL查询 - 项目: {project_name}, 类型: {query_type}")
    
    try:
        # 使用缓存的Prometheus连接
        prom = wechat_service.get_prometheus_client()
        if not prom:
            error_msg = "无法获取Prometheus客户端连接"
            logger.error(error_msg)
            return f"连接失败: {error_msg}"
        
        query_template = PROMQL_TEMPLATES.get(query_type) 
        if not query_template:
            error_msg = f"Unknown query type: {query_type}"
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        query = query_template.format(project_name=project_name)
        logger.info(f"构建的PromQL查询: {query}")
        
        # 执行查询
        query_res = prom.custom_query(query)
        execution_time = time.time() - start_time
        
        logger.info(f"PromQL查询完成 - 耗时: {execution_time:.2f}秒, 结果数量: {len(query_res) if query_res else 0}")
        
        if query_res:
            logger.debug(f"查询结果详情: {query_res}")
            msg = ''
            
            for i, result in enumerate(query_res):
                if i > 0:
                    msg += '\n' + '-' * 20 + '\n'
                
                query_node = result.get('metric', {}).get('node')
                query_name = result.get('metric', {}).get('__name__')
                query_instance = result.get('metric', {}).get('instance')
                query_pod = result.get('metric', {}).get('pod')
                query_phase = result.get('metric', {}).get('phase')
                raw_value = result.get('value', [None])[-1] if result.get('value') else None
                
                # 格式化数值，保留2位小数
                if raw_value is not None:
                    try:
                        # 尝试转换为浮点数并格式化
                        if isinstance(raw_value, (int, float)):
                            query_value = f"{float(raw_value):.2f}"
                        else:
                            # 如果是字符串，尝试转换
                            query_value = f"{float(raw_value):.2f}"
                    except (ValueError, TypeError):
                        # 如果转换失败，保持原值
                        query_value = str(raw_value)
                else:
                    query_value = None
                
                logger.debug(f"处理结果 {i+1}: 节点={query_node}, 指标={query_name}, 实例={query_instance}, 值={query_value}")
                
                lines = [
                    f'检测指标: {query_name}' if query_name else None,
                    f'检测实例: {query_instance}' if query_instance else None,
                    f'检测POD: {query_pod}' if query_pod else None,
                    f'检测POD状态: {query_phase}' if query_phase else None,
                    f'检测节点: {query_node}' if query_node else None,
                    f'检测值: {query_value}' if query_value is not None else None,
                ]
                msg += '\n'.join(line for line in lines if line is not None)
            
            logger.info(f"PromQL查询成功 - 返回数据长度: {len(msg)}字符")
            return msg
        else:
            no_data_msg = f"未找到查询项, 请确认 {project_name} 项目 是否监控 {query_type}"
            logger.warning(no_data_msg)
            return no_data_msg
            
    except Exception as e:
        execution_time = time.time() - start_time
        error_msg = f"执行PromQL查询异常: {e}"
        logger.error(f"{error_msg} - 耗时: {execution_time:.2f}秒")
        error_logger.error(f"PromQL查询失败 - 项目: {project_name}, 类型: {query_type}, 错误: {e}", exc_info=True)
        return f"查询执行失败: {str(e)}"

@app.route('/', methods=['GET', 'POST'])
def index():
    """首页"""
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return jsonify({
        'status': 'success',
        'message': f'Prometheus监控机器人服务运行中',
        'time': current_time,
        'supported_projects': len(projects),
        'supported_queries': len(promql_list)
    })

@app.route('/webhook', methods=['POST'])
def webhook():
    """企业微信Webhook接口"""
    try:
        data = request.json
        if not data:
            return jsonify({'error': '无效的请求数据'}), 400
        
        msg = data.get('msg', '').strip()
        if not msg:
            return jsonify({'error': '消息内容不能为空'}), 400
        
        logger.info(f"收到企业微信消息: {msg}")
        
        # 解析消息并查询Prometheus数据
        project_name, query_type = findproject_and_promql(msg)
        
        if not project_name or not query_type:
            # 返回帮助信息
            help_msg = generate_help_message()
            wechat_service.send_webhook_message(help_msg)
            return help_msg
        
        # 执行Prometheus查询
        result = build_promql(project_name, query_type)
        
        # 发送结果到企业微信
        formatted_result = wechat_service.format_prometheus_data(result, project_name, query_type)
        wechat_service.send_webhook_message(formatted_result, "markdown")
        
        return result
        
    except Exception as e:
        logger.error(f"处理企业微信消息异常: {e}")
        error_msg = f"处理请求时发生错误: {str(e)}"
        wechat_service.send_webhook_message(error_msg)
        return jsonify({'error': error_msg}), 500

@app.route('/wechat/callback', methods=['GET', 'POST'])
def wechat_callback():
    """企业微信应用回调接口"""
    try:
        if request.method == 'GET':
            # 获取验证参数
            msg_signature = request.args.get('msg_signature', '')
            timestamp = request.args.get('timestamp', '')
            nonce = request.args.get('nonce', '')
            echostr = request.args.get('echostr', '')
            
            logger.info(f"企业微信URL验证请求 - 签名: {msg_signature}, 时间戳: {timestamp}, 随机数: {nonce}")
            
            # 验证必要参数
            if not all([msg_signature, timestamp, nonce, echostr]):
                logger.error("缺少必要的验证参数")
                return "参数不完整", 400
            
            # 验证企业微信回调
            crypto = create_wechat_crypto()
            if crypto:
                try:
                    # 使用官方库验证并解密echostr
                    decrypted_echostr = crypto.check_signature(msg_signature, timestamp, nonce, echostr)
                    logger.info("企业微信URL验证成功")
                    return decrypted_echostr
                except InvalidSignatureException as e:
                    logger.error(f"企业微信URL验证失败: {e}")
                    return "签名验证失败", 400
            else:
                logger.error("企业微信加密器未正确配置")
                return "配置错误", 400
        
        elif request.method == 'POST':
            # 处理企业微信消息
            logger.info("收到企业微信POST回调请求")
            
            # 获取请求参数
            msg_signature = request.args.get('msg_signature', '')
            timestamp = request.args.get('timestamp', '')
            nonce = request.args.get('nonce', '')
            
            # 获取原始XML数据
            try:
                raw_data = request.get_data(as_text=True)
                logger.info(f"收到原始XML数据: {raw_data}")
                
                if not raw_data:
                    logger.error("请求体为空")
                    return jsonify({'error': '请求体为空'}), 400
                
            except Exception as e:
                logger.error(f"获取请求数据失败: {e}")
                return jsonify({'error': '获取请求数据失败'}), 400
            
            # 使用官方库解密消息
            try:
                crypto = create_wechat_crypto()
                if crypto:
                    # 使用官方库解密消息
                    decrypted_msg = crypto.decrypt_message(raw_data, msg_signature, timestamp, nonce)
                    logger.info(f"解密后的消息: {decrypted_msg}")
                    
                    # 解析XML消息
                    import xml.etree.ElementTree as ET
                    root = ET.fromstring(decrypted_msg)
                    
                    # 提取消息内容
                    msg_type = root.find('MsgType').text if root.find('MsgType') is not None else ''
                    content = root.find('Content').text if root.find('Content') is not None else ''
                    from_user = root.find('FromUserName').text if root.find('FromUserName') is not None else ''
                    
                    logger.info(f"解析消息 - 类型: {msg_type}, 内容: {content}, 发送者: {from_user}")
                    
                    # 处理文本消息
                    if msg_type == 'text' and content:
                        logger.info(f"收到文本消息: {content} from {from_user}")
                        
                        # 处理消息
                        project_name, query_type = findproject_and_promql(content)
                        
                        if project_name and query_type:
                            result = build_promql(project_name, query_type)
                            formatted_result = wechat_service.format_prometheus_data(result, project_name, query_type)
                            wechat_service.send_text_message(formatted_result, from_user)
                        else:
                            # 发送帮助信息
                            help_msg = generate_help_message()
                            wechat_service.send_text_message(help_msg, from_user)
                else:
                    logger.warning("未配置加密参数，无法解密消息")
                    return jsonify({'error': '未配置加密参数'}), 400
                    
            except Exception as e:
                logger.error(f"解密或解析消息失败: {e}")
                error_logger.error(f"消息处理失败: {e}", exc_info=True)
                return jsonify({'error': '消息处理失败'}), 400
            
            return jsonify({'status': 'success'})
    
    except Exception as e:
        logger.error(f"处理企业微信回调异常: {e}")
        error_logger.error(f"企业微信回调处理异常: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/wechat/robot', methods=['POST'])
def wechat_robot():
    """企业微信群聊机器人回调接口（用于接收@消息）"""
    try:
        data = request.json
        logger.info(f"收到群聊机器人消息: {data}")
        
        # 解析消息内容
        if data.get('msgtype') == 'text':
            content = data.get('text', {}).get('content', '').strip()
            from_user = data.get('from', {}).get('userid', '')
            chat_id = data.get('chatid', '')
            
            logger.info(f"收到群聊文本消息: {content} from {from_user} in chat {chat_id}")
            
            # 处理消息
            project_name, query_type = findproject_and_promql(content)
            
            if project_name and query_type:
                result = build_promql(project_name, query_type)
                formatted_result = wechat_service.format_prometheus_data(result, project_name, query_type)
                # 发送到群聊，@发送者
                wechat_service.send_group_message(formatted_result, [from_user])
            else:
                # 发送帮助信息
                help_msg = generate_help_message()
                wechat_service.send_group_message(help_msg, [from_user])
        
        return jsonify({'status': 'success'})
    
    except Exception as e:
        logger.error(f"处理群聊机器人消息异常: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/wechat/robot/webhook', methods=['POST'])
def wechat_robot_webhook():
    """企业微信群聊机器人Webhook接口（用于发送消息）"""
    try:
        data = request.json
        logger.info(f"收到群聊机器人Webhook消息: {data}")
        
        # 解析消息内容
        if data.get('msgtype') == 'text':
            content = data.get('text', {}).get('content', '').strip()
            from_user = data.get('from', {}).get('userid', '')
            chat_id = data.get('chatid', '')
            
            logger.info(f"收到群聊Webhook文本消息: {content} from {from_user} in chat {chat_id}")
            
            # 处理消息
            project_name, query_type = findproject_and_promql(content)
            
            if project_name and query_type:
                result = build_promql(project_name, query_type)
                formatted_result = wechat_service.format_prometheus_data(result, project_name, query_type)
                # 发送到群聊，@发送者
                wechat_service.send_group_message(formatted_result, [from_user])
            else:
                # 发送帮助信息
                help_msg = generate_help_message()
                wechat_service.send_group_message(help_msg, [from_user])
        
        return jsonify({'status': 'success'})
    
    except Exception as e:
        logger.error(f"处理群聊机器人Webhook消息异常: {e}")
        return jsonify({'error': str(e)}), 500


@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    try:
        # 检查Prometheus连接
        prom_response = requests.get(f'{PROMETHEUS_URL}/api/v1/query?query=up', timeout=5)
        prom_status = "healthy" if prom_response.status_code == 200 else "unhealthy"
        
        # 检查企业微信连接
        wechat_status = "healthy" if wechat_service.get_access_token() else "unhealthy"
        
        # 获取Prometheus连接缓存状态
        prom_connection_status = wechat_service.get_connection_status()
        
        return jsonify({
            'status': 'healthy',
            'prometheus': prom_status,
            'wechat': wechat_status,
            'prometheus_connection_cache': prom_connection_status,
            'timestamp': datetime.now().isoformat()
        })
    except Exception as e:
        logger.error(f"健康检查异常: {e}")
        return jsonify({
            'status': 'unhealthy',
            'error': str(e),
            'timestamp': datetime.now().isoformat()
        }), 500

if __name__ == '__main__':
    logger.info(f"启动Prometheus监控机器人服务 - {SERVER_HOST}:{SERVER_PORT}")
    
    # 启动连接清理线程
    cleanup_thread = threading.Thread(target=cleanup_connections_periodically, daemon=True)
    cleanup_thread.start()
    logger.info("Prometheus连接清理线程已启动")
    
    app.run(host=SERVER_HOST, port=SERVER_PORT, debug=False)
